# ''' Python TCP标定工具代码开始 '''

import numpy as np
import yaml
import os
import argparse
from scipy.spatial.transform import Rotation

class TCPCalibrator:
    """
    机器人工具中心点 (TCP) 标定器。

    该工具使用多点法，通过采集机器人工具以不同姿态接触
    空间中同一点时的法兰盘位姿，来计算 TCP 的位置。

    数学原理:
    对于第 i 次测量，我们有:
        P_ref = T_base_flange_i * P_flange_tcp
    其中:
        - P_ref 是空间固定参考点的齐次坐标 (未知, 但恒定)。
        - T_base_flange_i 是第 i 次测量时，机器人法兰在基坐标系下的位姿(4x4矩阵, 已知)。
        - P_flange_tcp 是待求的TCP在法兰坐标系下的齐次坐标 (未知, 恒定)。

    展开为旋转矩阵 R 和平移向量 t:
        p_ref = R_i @ p_tcp + t_i

    对于任意两次测量 i 和 j:
        R_i @ p_tcp + t_i = R_j @ p_tcp + t_j
    整理得:
        (R_i - R_j) @ p_tcp = t_j - t_i

    通过采集 N (N>=3) 组不同的姿态，我们可以构建一个超定线性方程组 Ax = b，
    并使用最小二乘法求解 p_tcp。
    """

    def __init__(self, config_path):
        """
        初始化TCP标定器。
        :param config_path: 配置文件的路径。
        """
        print("INFO: 初始化TCP标定器...")
        if not os.path.exists(config_path):
            raise FileNotFoundError(f"配置文件未找到: {config_path}")
            
        with open(config_path, 'r', encoding='utf-8') as f:
            self.config = yaml.safe_load(f)
        
        self.input_data_path = self.config['input_data_path']
        self.output_path = self.config['output_path']
        self.data_format = self.config.get('data_format', 'xyz_rxryrz').lower()
        self.units = self.config.get('units', 'mm')

        # 确保输出目录存在
        output_dir = os.path.dirname(self.output_path)
        if output_dir and not os.path.exists(output_dir):
            os.makedirs(output_dir)
            print(f"INFO: 已创建输出目录: {output_dir}")

    def _load_robot_poses(self):
        """
        从输入文件加载机器人法兰盘位姿数据。
        文件应为 .txt 或 .csv 格式，用逗号或空格分隔。
        """
        if not os.path.exists(self.input_data_path):
            raise FileNotFoundError(f"机器人位姿数据文件未找到: {self.input_data_path}")

        try:
            # 自动跳过注释行(以#开头)
            data = np.loadtxt(self.input_data_path, comments='#', delimiter=',')
            if data.ndim == 1: # 如果只有一行数据, 强制变为2D
                data = data.reshape(1, -1)
        except Exception as e:
            raise IOError(f"无法读取或解析数据文件 '{self.input_data_path}': {e}")

        poses = []
        num_cols = data.shape[1]

        print(f"INFO: 正在从 '{self.input_data_path}' 加载 {data.shape[0]} 个位姿点...")
        print(f"INFO: 数据格式被解析为 '{self.data_format}'。")

        for i, row in enumerate(data):
            try:
                t_vec = row[0:3]
                
                if self.data_format == 'xyz_rxryrz':
                    if num_cols != 6: raise ValueError(f"需要6列数据 (x,y,z,rx,ry,rz)，但找到 {num_cols} 列。")
                    rot_vec = row[3:6]
                    r = Rotation.from_rotvec(rot_vec)
                elif self.data_format == 'xyz_qwxyz':
                    if num_cols != 7: raise ValueError(f"需要7列数据 (x,y,z,qw,qx,qy,qz)，但找到 {num_cols} 列。")
                    # 注意: scipy的quaternion输入顺序是 [qx, qy, qz, qw]
                    quat = [row[4], row[5], row[6], row[3]]
                    r = Rotation.from_quat(quat)
                else:
                    raise ValueError(f"不支持的数据格式: '{self.data_format}'. 支持的格式: 'xyz_rxryrz', 'xyz_qwxyz'.")
                
                R_mat = r.as_matrix()
                poses.append({'R': R_mat, 't': t_vec})
            except Exception as e:
                raise ValueError(f"处理第 {i+1} 行数据时出错: {e}")

        print(f"INFO: 成功加载并转换了 {len(poses)} 个位姿。")
        return poses

    def calibrate(self):
        """
        执行TCP标定计算。
        """
        print("\nINFO: ============== 开始TCP标定 ==============")
        try:
            poses = self._load_robot_poses()
        except (FileNotFoundError, ValueError, IOError) as e:
            print(f"ERROR: 加载数据失败. {e}")
            return

        if len(poses) < 3:
            print(f"ERROR: 标定需要至少3个不同的位姿点，但只提供了 {len(poses)} 个。标定失败。")
            return

        # 使用第一个位姿作为参考，构建Ax=b方程组
        # (R_i - R_0) @ p_tcp = t_0 - t_i
        A = []
        b = []
        R0, t0 = poses[0]['R'], poses[0]['t']

        for i in range(1, len(poses)):
            Ri, ti = poses[i]['R'], poses[i]['t']
            A.append(Ri - R0)
            b.append(t0 - ti)
        
        # 将列表堆叠成一个大的矩阵和向量
        A_matrix = np.vstack(A)
        b_vector = np.concatenate(b)

        # 使用最小二乘法求解 p_tcp
        try:
            # result[0] 是解, result[1] 是残差, result[2] 是矩阵A的秩, result[3] 是A的奇异值
            p_tcp, residuals, rank, s = np.linalg.lstsq(A_matrix, b_vector, rcond=None)
            
            if rank < 3:
                 print("WARNING: 测量数据可能共线或共面，导致解不稳定。请尝试使用更多样化的姿态重新测量。")

            self._evaluate_and_save_results(p_tcp, poses, residuals)

        except np.linalg.LinAlgError as e:
            print(f"ERROR: 求解线性方程组时发生错误: {e}")
            print("ERROR: 请检查您的输入数据。确保机器人姿态有足够的变化。")

    def _evaluate_and_save_results(self, tcp_vector, poses, residuals):
        """
        评估标定结果的精度并保存。
        """
        print("\n----------- 标定结果 -----------")
        print(f"计算出的 TCP 坐标 (相对于法兰, 单位: {self.units}):")
        print(f"  x: {tcp_vector[0]:.4f}")
        print(f"  y: {tcp_vector[1]:.4f}")
        print(f"  z: {tcp_vector[2]:.4f}")
        print("---------------------------------\n")

        # --- 精度评估 ---
        # 基于计算出的TCP，反向计算每次测量时参考点在世界坐标系中的位置
        ref_points = []
        for pose in poses:
            p_ref = pose['R'] @ tcp_vector + pose['t']
            ref_points.append(p_ref)
        
        ref_points = np.array(ref_points)
        
        # 计算所有估计参考点的平均位置
        mean_ref_point = np.mean(ref_points, axis=0)

        # 计算每个点到平均位置的欧氏距离（误差）
        errors = np.linalg.norm(ref_points - mean_ref_point, axis=1)

        print("----------- 精度评估 -----------")
        print(f"估计的世界参考点平均位置: {mean_ref_point.round(4)}")
        print(f"标定一致性误差 (单位: {self.units}):")
        print(f"  - 平均误差: {np.mean(errors):.4f}")
        print(f"  - 最大误差: {np.max(errors):.4f}")
        print(f"  - 标准差:   {np.std(errors):.4f}")
        print(f"最小二乘法残差和: {residuals[0]:.4f}" if residuals.size > 0 else "N/A")
        print("---------------------------------\n")
        
        # --- 保存结果 ---
        data_to_save = {
            'tcp_vector': tcp_vector.tolist(),
            'units': self.units,
            'evaluation': {
                'mean_error': float(np.mean(errors)),
                'max_error': float(np.max(errors)),
                'std_dev': float(np.std(errors)),
                'lstsq_residuals': float(residuals[0]) if residuals.size > 0 else None,
                'estimated_world_point': mean_ref_point.tolist(),
            }
        }
        
        try:
            with open(self.output_path, 'w', encoding='utf-8') as f:
                yaml.dump(data_to_save, f, default_flow_style=False, sort_keys=False, indent=4)
            print(f"INFO: 标定结果已成功保存至: {self.output_path}")
        except Exception as e:
            print(f"ERROR: 保存结果到文件时失败: {e}")

        print("\n标定流程完成！")

# ===================================================================
# 主函数入口
# ===================================================================
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="机器人TCP标定工具.")
    parser.add_argument('--config', type=str, default='tcp_config.yaml', help='配置文件的路径')
    
    args = parser.parse_args()

    try:
        calibrator = TCPCalibrator(args.config)
        calibrator.calibrate()
    except (ValueError, FileNotFoundError, KeyError, TypeError, IOError) as e:
        print(f"\nFATAL ERROR: 处理失败. 详情: {e}")